################################################################################
#
# Copyright (c) 2007 Christopher J. Stawarz
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT.  IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
# BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.
#
################################################################################

from collections import defaultdict, deque
from functools import partial
from heapq import heappush, heappop, heapify
from types import GeneratorType
import os
import errno
import select
import sys
import time
import traceback

__author__ = 'Christopher Stawarz <cstawarz@csail.mit.edu>'
__version__ = '0.2.0'
__revision__ = int('$Revision: 5025 $'.split()[1])
__all__ = ["Timeout", "resume", "sleep", "readable", "writable",
          "read", "write", "accept", "recv", "send", "sendto",
          "queue", "taskmgr", "add", "run", "kill", 
          "FDReady", "FDAction", "ChildTask", "ResumeTask", 
          "SleepDelay", "YieldCondition", "QueueAction", "Queue"]


class Timeout(Exception):
    'Raised in a yielding task when an operation times out'
    pass

class ChildTask(object):
    def __init__(self, parent, task):
        self.parent = parent
        self.task = task

    def send(self, value):
        return self.task.send(value)

    def throw(self, type, value=None, traceback=None):
        return self.task.throw(type, value, traceback)

# NOT DONE YET: by lyricconch
# use yield resume(task) as GOTO
class ResumeTask(object):
    def __init__(self, task, important):
        self.task = task
        self.important = important

def resume(task, important):
    return ResumeTask(task=None, important=False)

class YieldCondition(object):
    """ Base class for objects that are yielded by a task to the task
    manager and specify the condition(s) under which the task should
    be restarted.  Only subclasses of this class are useful to
    application code. """

    def __init__(self, timeout=None):
        """ If timeout is None, the task will be suspended indefinitely
        until the condition is met.  Otherwise, if the condition is
        not met within timeout seconds, a Timeout exception will be
        raised in the yielding task."""
        self.task = None
        self.handle_expiration = None
        if timeout is None:
            self.expiration = None
        else:
            self.expiration = time.time() + float(timeout)

    @property
    def has_expires(self):
        return self.expiration is not None

class SleepDelay(YieldCondition):
    """ A task that yields the result of this function will be resumed
    after the specified number of seconds have elapsed. """

    def __init__(self, seconds):
        seconds = float(seconds)
        if seconds <= 0.0:
            raise ValueError("'seconds' must be greater than 0")
        super(SleepDelay, self).__init__(seconds)

def sleep(seconds):
    return SleepDelay(seconds)

class FDReady(YieldCondition):
    """ A task that yields an instance of this class will be suspended
    until a specified file descriptor is ready for I/O. """

    def __init__(self, fd, read=False, write=False, exc=False, timeout=None):
        """ Resume the yielding task when fd is ready for reading,
        writing, and/or "exceptional" condition handling.  fd can be
        any object accepted by select.select() (meaning an integer or
        an object with a fileno() method that returns an integer).
        Any exception raised by select() due to fd will be re-raised
        in the yielding task.
        If timeout is not None, a Timeout exception will be raised in
        the yielding task if fd is not ready after timeout seconds
        have elapsed. """

        super(FDReady, self).__init__(timeout)
        self.fd = fd if isinstance(fd, int) else fd.fileno()
        self.io = (read, write, exc)
        if not any(self.io):
            raise ValueError("'read', 'write', and 'exc' cannot all be false")

    def __call__(self):
        return None

    def fileno(self):
        'Return the file descriptor on which the yielding task is waiting'
        return self.fd

    def _enter_select(self, *fdsets):
        for enter, fdset in zip(self.io, fdsets):
            if enter: fdset.add(self)

    def _leave_select(self, *fdsets):
        for leave, fdset in zip(self.io, fdsets):
            if leave: fdset.discard(self)

def readable(fd, timeout=None):
    return FDReady(fd, read=True, timeout=timeout)

def writable(fd, timeout=None):
    return FDReady(fd, write=True, timeout=timeout)

class FDAction(FDReady):
    """ A task that yields an instance of this class will be suspended
    until an I/O operation on a specified file descriptor is complete. """

    def __init__(self, fd, func, args=(), kwargs={}, read=False, write=False, exc=False):
        """ Resume the yielding task when fd is ready for reading,
        writing, and/or "exceptional" condition handling.  fd can be
        any object accepted by select.select() (meaning an integer or
        an object with a fileno() method that returns an integer).
        Any exception raised by select() due to fd will be re-raised
        in the yielding task.

        The value of the yield expression will be the result of
        calling func with the specified args and kwargs (which
        presumably performs a read, write, or other I/O operation on
        fd).  If func raises an exception, it will be re-raised in the
        yielding task.  Thus, FDAction is really just a convenient
        subclass of FDReady that requests that the task manager
        perform an I/O operation on the calling task's behalf.

        If kwargs contains a timeout argument that is not None, a
        Timeout exception will be raised in the yielding task if fd is
        not ready after timeout seconds have elapsed. """

        timeout = kwargs.pop('timeout', None)
        super(FDAction, self).__init__(fd, read, write, exc, timeout)
        self.func = func
        self.args = args
        self.kwargs = kwargs

    def __call__(self):
        return self.func(*(self.args), **(self.kwargs))
    
def read(fd, *args, **kwargs):
    func = partial(os.read, fd) if isinstance(fd, int) else fd.read
    return FDAction(fd, func, args, kwargs, read=True)

def readline(fd, *args, **kwargs):
    return FDAction(fd, fd.readline, args, kwargs, read=True)

def write(fd, *args, **kwargs):
    func = (partial(os.write, fd) if isinstance(fd, int) else fd.write)
    return FDAction(fd, func, args, kwargs, write=True)

def accept(sock, *args, **kwargs):
    return FDAction(sock, sock.accept, args, kwargs, read=True)

def recv(sock, *args, **kwargs):
    return FDAction(sock, sock.recv, args, kwargs, read=True)

def recvfrom(sock, *args, **kwargs):
    return FDAction(sock, sock.recvfrom, args, kwargs, read=True)

def send(sock, *args, **kwargs):
    return FDAction(sock, sock.send, args, kwargs, write=True)

def sendto(sock, *args, **kwargs):
    return FDAction(sock, sock.sendto, args, kwargs, write=True)

class Queue(object):
    """ A multi-producer, multi-consumer FIFO queue (similar to
    Queue.Queue) that can be used for exchanging data between tasks """

    def __init__(self, contents=(), maxsize=0):
        """ Create a new Queue instance.  contents is a sequence (empty by
        default) containing the initial contents of the queue.  If
        maxsize is greater than 0, the queue will hold a maximum of
        maxsize items, and put() will block until space is available
        in the queue. """

        self.maxsize = int(maxsize)
        self._queue = deque(contents)

    def __len__(self):
        return len(self._queue)

    def _get(self):
        return self._queue.popleft()

    def _put(self, item):
        self._queue.append(item)

    def empty(self):
        return len(self) == 0

    def full(self):
        return 0 < self.maxsize <= len(self)

    def get(self, timeout=None):
        return QueueAction(self, timeout=timeout)

    def put(self, item, timeout=None):
        return QueueAction(self, item, timeout=timeout)

def queue(contents=(), maxsize=0):
    return Queue(contents, maxsize)

class QueueAction(YieldCondition):
    NO_ITEM = object()

    def __init__(self, queue, item=NO_ITEM, timeout=None):
        super(QueueAction, self).__init__(timeout)
        if not isinstance(queue, Queue):
            raise TypeError("'queue' must be a Queue instance")
        self.queue = queue
        self.item = item

class TaskManager(object):
    """ Engine for running a set of cooperatively-multitasking tasks
    within a single Python thread """

    def __init__(self):
        """  Create a new TaskManager instance.  Generally, there will only
        be one of these per Python process.  If you want to run two
        existing instances simultaneously, merge them first, then run
        one or the other. """
        self._queue = deque()
        self._io_waits = [set(), set(), set()]
        self._queue_waits = defaultdict(lambda: (deque(), deque()))
        self._timeouts = []
        self._suspused = []

    def merge(self, other=None):
        """  Merge this TaskManager with another.  After the merge, the two
        objects share the same (merged) internal data structures, so
        either can be used to manage the combined task set."""
        if other is None: other = taskmgr
        assert isinstance(other, TaskManager), "must be a TaskManager instance"

        self._queue.extend(other._queue)
        self._io_waits[0] |= other._io_waits[0]
        self._io_waits[1] |= other._io_waits[1]
        self._io_waits[2] |= other._io_waits[2]
        self._queue_waits.update(other._queue_waits)
        self._timeouts.extend(other._timeouts)
        heapify(self._timeouts)

        other._queue = self._queue
        other._io_waits = self._io_waits
        other._queue_waits = self._queue_waits
        other._timeouts = self._timeouts

    def add(self, task):
        'Add a new task (i.e. a generator instance) to the run queue'
        assert isinstance(task, GeneratorType)
        self._enqueue(task)
        #self._weakref[task]= None
    
    # TODO: 
    def kill(self, task):
        raise NotImplemented
        for i in xrange(self._queue):
            if i[0] is task: self._queue.remove(i)
            if isinstance(i[0], ChildTask): pass

    def _enqueue(self, task, input=None, exc_info=()):
        self._queue.append((task, input, exc_info))

    def run(self, interval=None):
        try:
            while self.has_runnable or self.has_io_waits or self.has_timeouts:
                self.run_next(interval)
        except:
            traceback.print_exc()
            sys.stderr.flush()

    @property
    def has_runnable(self):
        return bool(self._queue)
    @property
    def has_io_waits(self):
        return any(self._io_waits)
    @property
    def has_timeouts(self):
        return bool(self._timeouts)

    def run_next(self, timeout=None):
        """  Perform one iteration of the run cycle: check whether any
        pending I/O operations can be performed, check whether any
        timeouts have expired, then run all currently runnable tasks.

        The timeout argument specifies the maximum time to wait for
        some task to become runnable.  If timeout is None and there
        are no currently runnable tasks, but there are tasks waiting
        to perform I/O or time out, then this method will block until
        at least one of the waiting tasks becomes runnable.  To
        prevent this method from blocking indefinitely, use timeout to
        specify the maximum number of seconds to wait.

        If there are runnable tasks in the queue when run_next() is
        called, then it will check for I/O readiness using a
        non-blocking call to select() (i.e. a poll), and only
        already-expired timeouts will be handled.  This ensures both
        that the task manager is never idle when tasks can be run and
        that tasks waiting for I/O never starve. """

        if self.has_io_waits:
            self._handle_io_waits(self._fix_run_timeout(timeout))

        if self.has_timeouts:
            self._handle_timeouts(self._fix_run_timeout(timeout))

        # Run all tasks currently in the queue
        for dummy in xrange(len(self._queue)):
            task, input, exc_info = self._queue.popleft()
            try:
                output = task.throw(*exc_info)if exc_info else task.send(input)
            except StopIteration, e:
                if isinstance(task, ChildTask):
                    output = None if not e.args else e.args if len(e.args) > 1 else e.args[0]
                    self._enqueue(task.parent, input=output)
            except:
                if not isinstance(task, ChildTask): raise
                self._enqueue(task.parent, exc_info=sys.exc_info())
            else:
                self._handle_task_output(task, output)

    def _fix_run_timeout(self, timeout):
        if self.has_runnable: return 0.0
        if not self.has_timeouts: return timeout
        expect = max(0, self._timeouts[0][0] - time.time())
        return expect if timeout is None else min(timeout, expect)

    def _handle_io_waits(self, timeout):
        try:
            read, write, exc = select.select(*self._io_waits + [timeout])
        except (TypeError, ValueError):
            self._remove_bad_file_descriptors()
        except (select.error, IOError), err:
            if err.args[0] == errno.EINTR:
                pass
            elif err.args[0] in (errno.EBADF, errno.WSAENOTSOCK):
                self._remove_bad_file_descriptors()
            else:
                raise
        else:
            for fd in read + write + exc:
                try:
                    input = fd()
                    if input is Ellipsis: continue
                    self._enqueue(fd.task, input=input)
                except:
                    self._enqueue(fd.task, exc_info=sys.exc_info())
                fd._leave_select(*self._io_waits)
                if fd.has_expires:
                    self._remove_timeout(fd)

    def _remove_bad_file_descriptors(self):
        for fd in set.union(*self._io_waits):
            try:
                select.select([fd], [fd], [fd], 0.0)
            except:
                self._enqueue(fd.task, exc_info=sys.exc_info())
                fd._leave_select(*self._io_waits)
                if fd.has_expires:
                    self._remove_timeout(fd)

    def _add_timeout(self, item, handler):
        item.handle_expiration = handler
        heappush(self._timeouts, (item.expiration, item))

    def _remove_timeout(self, item):
        self._timeouts.remove((item.expiration, item))
        heapify(self._timeouts)

    def _handle_timeouts(self, timeout):
        if timeout > 0.0:
            time.sleep(timeout)
        now = time.time()
        while self._timeouts and self._timeouts[0][0] <= now:
            item = heappop(self._timeouts)[1]
            if isinstance(item, SleepDelay):
                self._enqueue(item.task)
            else:
                self._enqueue(item.task, exc_info=(Timeout,))
                item.handle_expiration()

    def _handle_task_output(self, task, output):
        if isinstance(output, GeneratorType):
            self._enqueue(ChildTask(task, output))
        elif isinstance(output, YieldCondition):
            output.task = task
            if isinstance(output, SleepDelay):
                self._add_timeout(output, None)
            elif isinstance(output, FDReady):
                self._handle_fdready(task, output)
            elif isinstance(output, QueueAction):
                self._handle_queue_action(task, output)
            elif isinstance(output, ResumeTask):
                self._handle_resume_task(task, output)
        else:
            self._enqueue(task, input=output)

    # TODO: 
    def _handle_resume_task(self, task, output):
        # yield resume(task) support
        target = output.target
        if target in self._suspused:
            self._suspused.append(task)
            self._suspused.remove(target)
            if output.important:
                self._queue.appendleft((target, task, None))
            else:
                self._queue.append((target, task, None))
        elif target is None:
            self._suspused.append(task)
        else:
            self._enqueue(task)

    def _handle_fdready(self, task, output):
        output._enter_select(*self._io_waits)
        if output.has_expires:
            self._add_timeout(output, lambda: output._leave_select(*self._io_waits))

    def _handle_queue_action(self, task, output):
        get_waits, put_waits = self._queue_waits[output.queue]

        if output.item is output.NO_ITEM:
            # Action is a get
            if output.queue.empty():
                get_waits.append(output)
                if output.has_expires:
                    self._add_timeout(output, lambda: get_waits.remove(output))
            else:
                item = output.queue._get()
                self._enqueue(task, input=item)
                if put_waits:
                    action = put_waits.popleft()
                    output.queue._put(action.item)
                    self._enqueue(action.task)
                    if action.has_expires:
                        self._remove_timeout(action)
        else:
            # Action is a put
            if output.queue.full():
                put_waits.append(output)
                if output.has_expires:
                    self._add_timeout(output, lambda: put_waits.remove(output))
            else:
                output.queue._put(output.item)
                self._enqueue(task)
                if get_waits:
                    action = get_waits.popleft()
                    item = output.queue._get()
                    self._enqueue(action.task, input=item)
                    if action.has_expires:
                        self._remove_timeout(action)

taskmgr = TaskManager()

def add(task):
    taskmgr.add(task)

def kill(task):
    taskmgr.kill(task)

def run(interval=0.05):
    taskmgr.run(interval)

Return = StopIteration

if __name__ == '__main__':

    def printer(name):
        for i in xrange(1, 4):
            print '%s:\t%d' % (name, i)
            yield

    t = TaskManager()
    t.add(printer('first'))
    t.add(printer('second'))
    t.add(printer('third'))

    queue = Queue()

    def receiver():
        print 'receiver started'
        print 'receiver received: %s' % (yield queue.get())
        print 'receiver finished'

    def sender():
        print 'sender started'
        yield queue.put('from sender')
        print 'sender finished'

    def bad_descriptor():
        print 'bad_descriptor running'
        try:
            yield readable(12)
        except:
            print 'exception in bad_descriptor:', sys.exc_info()[1]

    def sleeper():
        print 'sleeper started'
        yield sleep(1)
        print 'sleeper finished'

    def timeout_immediately():
        print 'timeout_immediately running'
        try:
            yield Queue().get(timeout=0)
        except Timeout:
            print 'timeout_immediately timed out'

    t2 = TaskManager()
    t2.add(receiver())
    t2.add(bad_descriptor())
    t2.add(sender())
    t2.add(sleeper())
    t2.add(timeout_immediately())

    def parent():
        print 'child returned: ' + str((yield child()))
        try:
            yield child(raise_exc=True)
        except:
            print 'exception in child:', sys.exc_info()[1]

    def child(raise_exc=False):
        yield
        if raise_exc:
            raise RuntimeError('foo')
        raise StopIteration(1, 2, 3)

    t3 = TaskManager()
    t3.add(parent())

    t.merge(t2)
    t.merge(t3)
    t.run()

    assert not(t.has_runnable or t.has_io_waits or t.has_timeouts)
